"""
任务管理服务
处理任务创建、状态跟踪、重试逻辑
"""
import time
from typing import Dict, Any, Optional, List
from datetime import datetime
from tenacity import retry, stop_after_attempt, wait_exponential
from loguru import logger

from storyforge.models import Job
from storyforge.utils import generate_id, log_job, save_json, load_json


class JobManager:
    """任务管理器"""
    
    def __init__(self, jobs_dir: str = "./logs"):
        """
        初始化任务管理器
        
        Args:
            jobs_dir: 任务数据目录
        """
        self.jobs_dir = jobs_dir
        self.jobs: Dict[str, Job] = {}
        self._load_jobs()
    
    def _load_jobs(self):
        """
        从文件加载任务数据（如果存在）
        
        从 jobs.json 文件中读取并恢复任务状态到内存
        """
        jobs_file = f"{self.jobs_dir}/jobs.json"
        data = load_json(jobs_file)
        if data:
            for job_id, job_data in data.items():
                try:
                    self.jobs[job_id] = Job(**job_data)
                except Exception as e:
                    logger.warning(f"加载任务失败 {job_id}: {e}")
    
    def _save_jobs(self):
        """
        保存任务数据到文件
        
        将内存中的任务状态持久化到 jobs.json 文件
        """
        jobs_file = f"{self.jobs_dir}/jobs.json"
        data = {job_id: job.dict() for job_id, job in self.jobs.items()}
        save_json(data, jobs_file)
    
    def create_job(
        self,
        job_type: str,
        inputs: Dict[str, Any],
        external_provider: Optional[str] = None,
        external_job_id: Optional[str] = None
    ) -> Job:
        """
        创建新任务
        
        Args:
            job_type: 任务类型
            inputs: 输入参数
            external_provider: 外部提供商
            external_job_id: 外部任务ID
            
        Returns:
            创建的Job对象
        """
        job_id = generate_id("job_")
        job = Job(
            id=job_id,
            type=job_type,
            inputs=inputs,
            status="pending",
            externalProvider=external_provider,
            externalJobId=external_job_id
        )
        
        self.jobs[job_id] = job
        self._save_jobs()
        log_job(job.dict())
        
        logger.info(f"创建任务: {job_id} ({job_type})")
        return job
    
    def get_job(self, job_id: str) -> Optional[Job]:
        """
        获取任务
        
        Args:
            job_id: 任务ID
            
        Returns:
            Job对象，不存在返回None
        """
        return self.jobs.get(job_id)
    
    def update_job(
        self,
        job_id: str,
        status: Optional[str] = None,
        outputs: Optional[Dict[str, Any]] = None,
        error: Optional[str] = None,
        cost: Optional[float] = None
    ) -> bool:
        """
        更新任务状态
        
        Args:
            job_id: 任务ID
            status: 新状态
            outputs: 输出结果
            error: 错误信息
            cost: 成本
            
        Returns:
            是否成功
        """
        job = self.jobs.get(job_id)
        if not job:
            logger.warning(f"任务不存在: {job_id}")
            return False
        
        if status:
            job.status = status
        if outputs is not None:
            job.outputs = outputs
        if error:
            job.error = error
        if cost is not None:
            job.cost = cost
        
        job.updatedAt = datetime.now()
        self._save_jobs()
        log_job(job.dict())
        
        logger.debug(f"更新任务: {job_id} -> {status}")
        return True
    
    def retry_job(self, job_id: str) -> bool:
        """
        重试任务
        
        Args:
            job_id: 任务ID
            
        Returns:
            是否成功
        """
        job = self.jobs.get(job_id)
        if not job:
            return False
        
        job.retry += 1
        job.status = "pending"
        job.error = None
        job.updatedAt = datetime.now()
        
        self._save_jobs()
        log_job(job.dict())
        
        logger.info(f"重试任务: {job_id} (第{job.retry}次)")
        return True
    
    def list_jobs(self, status: Optional[str] = None) -> List[Job]:
        """
        列出任务
        
        Args:
            status: 筛选状态，None表示全部
            
        Returns:
            任务列表
        """
        if status:
            return [job for job in self.jobs.values() if job.status == status]
        return list(self.jobs.values())


# 全局任务管理器实例
_job_manager: Optional[JobManager] = None


def get_job_manager(jobs_dir: str = "./logs") -> JobManager:
    """
    获取任务管理器实例（单例）
    
    Args:
        jobs_dir: 任务数据目录
        
    Returns:
        JobManager实例
    """
    global _job_manager
    if _job_manager is None:
        _job_manager = JobManager(jobs_dir)
    return _job_manager

